Lr scheduler

sd-scriptsにおけるLr scheduler

概要

• 学習率を時系列変化させるLearning rate専用のスケジューラ。
  学習のstep進行に応じて、学習率Learning rateへ0～1の値を掛けて学習率を調整する。1を基準として低下させていく。
  スケジューラによって、Lrの低下具合が異なる。

目的

• Lr schedulerの一番の目的は、過学習しないようにLrに徐々にブレーキを掛けること。
  • 「学習したいものを絞り出す能力」ともいえる。
    • 通常、適度に大きなLrは構図やポーズを学習するのには便利だが、Lrスケジューラで調整しない限り、土台である構図が安定せず、なかなかディティール学習が進まない。ディティールが安定しないと、構図の学習も疎かになるという悪循環に陥りがちである
    • 1stepで１つのLrしか適用できない学習の仕様上、高いLrで学習→クール期間を挟みつつ小さいLrで学習する方法が最も安定する*1
  • Lrを落とす結果、内部的には学習量は落ちるが、ウェイトが増えたり減ったりの無駄な計算が減るため、高めのLrと組み合わあわせて適切に設定することで結果として学習時間の短縮に繋がる。

使用方法

• 使用するには、次の行を編集することでオプション指定を行います。

     --lr_scheduler=

• 選択したlr_schedulerによっては追加のオプションが必要な場合があります。
• オプティマイザーがAdaFactorかつ、relative_step=trueの場合、Lr schedulerは自動的にadafactorが選択される。
• sd-scriptsのプリセット以外にカスタムスケジューラも使用可能。

Lr Schedulerの種類

sd-scriptsのプリセットは下表の通り。

赤：constant、黄：linear、青：cosine
※イメージ図なので、実際の値とは多少差異有り。

Lr scheduler選定の留意点

基本的にcosine_with_restartsで多くのニーズを満たせるはずだが、
オプティマイザーやEpoch数（全step中に教師画像が何回・どんな分布で登場するか）の相性の都合で、他のLr scehdulerを選択する場合がある。

基本的な考え方

一概に答えはありませんが、下記の傾向はあります。

• Lrの時間積分値が大きいほど、全体の学習量が増える。→Epoch数を減らせる（時短）。
• 自動型よりも手動型オプティマイザー（Lr固定）の方が、スケジューラ選定が比較的重要。なぜなら、下記要件を同時に満たすには、スケジューラ以外に頼れないからです。
  • Epoch前半：Lr倍率が大きい方が良い。学習量が増え、局所最適解に留まらずに済む
    （目標とは程遠いのになかなか学習が進まない状態のこと。谷を越えるとも言う）
  • Epoch後半：Lr倍率は適度に小さい方が良い。過学習せずにディティールを学習できる。

オプティマイザーとの相性

オプティマイザーによっては、その専用のパラメータが前stepの状態を参照している場合があります。
Lionが該当。不安定なスケジューラは不安定なウェイト更新をもたらし、間接的にオプティマイザーの状態に影響を与えるでしょう。

lr_warmup_steps併用も考慮する。

• 学習初期のstepではオプティマイザーの状態が安定していないため、warmupにより徐々に学習率を倍率１まで上げる方法がある。

  --lr_warmup_steps=10
  #この場合は0～10stepの間、0からLrの設定値まで直線的に増加する。

• 学習開始直後に不安定になる場合に適用すると良い。
  • もし目安が欲しいなら、lossを監視して急激なloss変動が落ち着くあたりのstep数を指定すれば良いと思う。

Tips

cosine系の欠点と対策

• cosineの欠点
  cosine系は、後半のstep全体のかなり長いstep区間で学習率がほぼ0になる。
  そのため、教師画像の何割かが無駄になるおそれがある
  • たとえば、1epoch（かつrestart数=1）にしてしまうと、まったく学習されない画像が出てくる。
  • 特に、有効なLearning rateのバンドが狭いオプティマイザーを使用するほど、ほとんどの学習が無駄か、特定の教師画像に偏った学習になる。
  • 自動型オプティマイザーの場合、cosineのせいでLrが微妙なところで安定してしまうため。cosineとは相性が悪い（本来高性能なprodigyが使いにくい一因である。）
• 対策
  Lr schedulerを改造して底上げを図ると良い。
  • cosine_with_restartsの場合の改造例は下記の通り。
    • optimization.py*5を参照し、下記のdef lr_lambda部分を修正。
      min_value=0.5として、最小値が0.5までしか下がらないようにする。min_valueの値は0.3～0.6程度の範囲で調整すると良い。
      コード改造例

      ※わからないことがあれば、ChatGPTに聞けば答えてくれるはず。

      
      def get_cosine_with_hard_restarts_schedule_with_warmup(
          optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
      ) -> LambdaLR:
          def lr_lambda(current_step):
              if current_step < num_warmup_steps:
                  return float(current_step) / float(max(1, num_warmup_steps))
              progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
              if progress >= 1.0:
                  return 0.0
      
              # Learning rateの最小値と最大値を設定する。
              max_value=1.0
              min_value=0.5
      
              # 既存の数式を、min_valueを考慮した数式に変換する。
              return max(0.0, min_value + (max_value - min_value) * (0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0)))))
          return LambdaLR(optimizer, lr_lambda, last_epoch)

• 対策後の状態

  Lr_scheduler係数の値。最低値が0.5まで増加した。その結果、学習がみるみる進む。
  緑が変更前、紫が変更後。横軸はstep数
  
  
  下図はProdigyによる補正後Lr*6。上図のグラフと同一のデータ。
  cosineの影響で学習率が十分に大きくならないまま安定し、再びLr scheduler=１に戻っても低い学習率を維持している。
  

  • 今回の検証では、restart間隔がProdigyのLr安定よりも長いため、低い学習率で安定したようだが、それ以外の状況だとどうなるかわからない。いずれにせよ、Lr schedulerを底上げした方が、schedulerの支配を受けにくい安定した結果が得られるだろう。

その他

使用者のコメントです。